{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "still-consistency",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import numpy as np\n",
    "from tensorflow.keras.datasets import mnist\n",
    "\n",
    "from dip import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "blessed-pizza",
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train, x_test = x_train / 255.0, x_test / 255.0\n",
    "x_train = x_train.reshape(-1, 1, 28, 28)\n",
    "x_test = x_test.reshape(-1, 1, 28, 28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "virtual-listing",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NN(Model):\n",
    "    def __init__(self):\n",
    "        self.conv1_W = Tensor.uniform(32, 1, 5, 5)\n",
    "        self.conv1_b = Tensor.uniform(1, 32, 1, 1)\n",
    "        self.conv2_W = Tensor.uniform(64, 32, 3, 3)\n",
    "        self.conv2_b = Tensor.uniform(1, 64, 1, 1)\n",
    "        self.W1 = Tensor.uniform(1600, 128)\n",
    "        self.b1 = Tensor.zeros(1, 128)\n",
    "        self.W2 = Tensor.uniform(128, 10)\n",
    "        self.b2 = Tensor.zeros(1, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = (x.conv2d(self.conv1_W)+self.conv1_b).relu()\n",
    "        x = x.max_pool2d().dropout()\n",
    "        x = (x.conv2d(self.conv2_W)+self.conv2_b).relu()\n",
    "        x = x.max_pool2d().dropout()\n",
    "        x = x.reshape((-1, 1600))\n",
    "        x = (x@self.W1+self.b1).relu().dropout()\n",
    "        x = (x@self.W2+self.b2).logsoftmax()\n",
    "        return x\n",
    "\n",
    "    def load(self, fn='vars.json'):\n",
    "        with open(fn, 'r') as f:\n",
    "            params = json.load(f)\n",
    "        cW1, cb1 = params['conv1']\n",
    "        self.conv1_W = Tensor(np.array(cW1))\n",
    "        self.conv1_b = Tensor(np.array(cb1).reshape(1, -1, 1, 1))\n",
    "        cW2, cb2 = params['conv2']\n",
    "        self.conv2_W = Tensor(np.array(cW2))\n",
    "        self.conv2_b = Tensor(np.array(cb2).reshape(1, -1, 1, 1))\n",
    "        W1, b1 = params['fc1']\n",
    "        self.W1 = Tensor(np.array(W1))\n",
    "        self.b1 = Tensor(np.array(b1).reshape(1, -1))\n",
    "        W2, b2 = params['fc2']\n",
    "        self.W2 = Tensor(np.array(W2))\n",
    "        self.b2 = Tensor(np.array(b2).reshape(1, -1))\n",
    "\n",
    "    def save(self, fn='vars.json'):\n",
    "        params = {}\n",
    "        params['conv1'] = (self.conv1_W.data.tolist(),\n",
    "                           self.conv1_b.data.tolist())\n",
    "        params['conv2'] = (self.conv2_W.data.tolist(),\n",
    "                           self.conv2_b.data.tolist())\n",
    "        params['fc1'] = (self.W1.data.tolist(), self.b1.data.tolist())\n",
    "        params['fc2'] = (self.W2.data.tolist(), self.b2.data.tolist())\n",
    "        with open(fn, 'w') as f:\n",
    "            json.dump(params, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "broad-reward",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = NN()\n",
    "m.load()\n",
    "opt = Adam(m.get_params())\n",
    "def lossfn(o, y): return (-o*Tensor(np.eye(10)[y])).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "loss -1.66187 accuracy 99.21875: 100%|██████████| 1/1 [04:36<00:00, 276.72s/it]\n"
     ]
    }
   ],
   "source": [
    "loss, acc = train(m, x_train, y_train, opt, 1, 128, lossfn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Acc: 99.369%: 100%|██████████| 40/40 [00:17<00:00,  2.32it/s]0.9937\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(evaluate(m, x_test, y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.save()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}